# -*- coding: utf-8 -*- 
# @Time : 2021/3/26/026 10:29 
# @Author : BuildRT
# @File : question_classifier.py
import os
import ahocorasick

class QuestionClassifier:
    def __init__(self):
        cur_dir = '/'.join(os.path.abspath(__file__).split('/')[:-1])
        cur_dir = os.path.dirname(__file__)
        # 　特征词路径
        self.groups_path = os.path.join(cur_dir, 'dict/groups.txt')
        self.foodCategory_path = os.path.join(cur_dir, 'dict/food_category.txt')
        self.food_path = os.path.join(cur_dir, 'dict/food.txt')
        self.nutrientDetail_path = os.path.join(cur_dir, 'dict/nutrients_detail.txt')

        # 加载特征词
        self.groups_wds = [i.strip() for i in open(self.groups_path,encoding="utf-8") if i.strip()]
        self.foodCategory_wds = [i.strip() for i in open(self.foodCategory_path, encoding="utf-8") if i.strip()]
        self.food_wds = [i.strip() for i in open(self.food_path, encoding="utf-8") if i.strip()]
        self.nutrientDetail_wds = [i.strip() for i in open(self.nutrientDetail_path, encoding="utf-8") if i.strip()]

        self.region_words = set(self.groups_wds + self.foodCategory_wds + self.food_wds + self.nutrientDetail_wds)

        # 构造领域actree
        self.region_tree = self.build_actree(list(self.region_words))
        # 构建词典
        self.wdtype_dict = self.build_wdtype_dict()

        # 问句疑问词
        # 询问大类食物下有哪些小类
        self.groupsFoodCategory_qwds = ['包含']
        # 询问某个小类是属于哪个大类的
        self.foodCategoryGroups_qwds = ['大类']
        # 询问某个小类中有哪些详细的食物
        self.foodCategoryFood_qwds = ['食物']
        # 询问某个食物是属于哪个小类的
        self.foodFoodCategory_qwds = ['种类']
        # 询问某个食物的能量与相关成分
        self.foodEnergy_qwds = ['能量', '能量与相关成分']
        # 询问某个食品有哪些维生素
        self.foodVitamin_qwds = ['维生素', '维他命']
        # 询问某个食品中的矿物质
        self.foodMineral_qwds = ['矿物质']
        # 询问某个食品中的脂肪酸
        self.foodFattyAcid_qwds = ['脂肪酸']

        print('model init finished ......')

        return

    '''分类主函数'''
    def classify(self, question):
        data = {}
        food_dict = self.check_food(question)
        if not food_dict:
            return {}
        data['args'] = food_dict
        # 收集问句当中所涉及到的实体类型
        types = []
        for type_ in food_dict.values():
            types += type_
        question_type = 'others'

        question_types = []

        # 某个大类中有哪些小类
        if self.check_words(self.groupsFoodCategory_qwds, question) and ('groups' in types):
            question_type = 'groups_foodCategory'
            question_types.append(question_type)

        if self.check_words(self.foodCategoryGroups_qwds, question) and ('foodCategory' in types):
            question_type = 'foodCategory_groups'
            question_types.append(question_type)

        # 某个小类中有哪些具体的食物
        if self.check_words(self.foodCategoryFood_qwds, question) and ('foodCategory' in types):
            question_type = 'foodCategory_food'
            question_types.append(question_type)

        if self.check_words(self.foodFoodCategory_qwds, question) and ('food' in types):
            question_type = 'food_foodCategory'
            question_types.append(question_type)

        # 某个食物中的能量与相关成分
        if self.check_words(self.foodEnergy_qwds, question) and ('food' in types):
            question_type = 'food_energy'
            question_types.append(question_type)

        # 某个食物中的维生素
        if self.check_words(self.foodVitamin_qwds, question) and ('food' in types):
            question_type = 'food_vitamin'
            question_types.append(question_type)

        # 某个食物中的矿物质
        if self.check_words(self.foodMineral_qwds, question) and ('food' in types):
            question_type = 'food_mineral'
            question_types.append(question_type)

        # 某个食物中的脂肪酸
        if self.check_words(self.foodFattyAcid_qwds, question) and ('food' in types):
            question_type = 'food_fatty_acid'
            question_types.append(question_type)

        # 将多个分类结果进行合并处理，组装成一个字典
        data['question_types'] = question_types
        return data


    '''构造词对应的类型'''
    def build_wdtype_dict(self):
        wd_dict = dict()
        for wd in self.region_words:
            wd_dict[wd] = []
            if wd in self.groups_wds:
                wd_dict[wd].append('groups')
            if wd in self.foodCategory_wds:
                wd_dict[wd].append('foodCategory')
            if wd in self.food_wds:
                wd_dict[wd].append('food')
            if wd in self.nutrientDetail_wds:
                wd_dict[wd].append('nutrientDetail')
        return wd_dict

    '''构造actree，加速过滤'''
    def build_actree(self, wordlist):
        actree = ahocorasick.Automaton()
        for index, word in enumerate(wordlist):
            actree.add_word(word, (index, word))
        actree.make_automaton()
        return actree

    '''问句过滤'''
    def check_food(self, question):
        region_wds = []
        for i in self.region_tree.iter(question):   # 看ac树节点上有没有这个词
            wd = i[1][1]
            region_wds.append(wd)
        stop_wds = []
        # 过滤
        for wd1 in region_wds:
            for wd2 in region_wds:
                if wd1 in wd2 and wd1 != wd2:
                    stop_wds.append(wd1)
        final_wds = [i for i in region_wds if i not in stop_wds]
        # 当前词语的类型
        final_dict = {i:self.wdtype_dict.get(i) for i in final_wds}

        return final_dict

    '''基于特征词进行分类'''
    def check_words(self, wds, sent):  # 判断关键词是否在这句话中
        for wd in wds:
            if wd in sent:
                return True
        return False

if __name__ == '__main__':
    handler = QuestionClassifier()
    while 1:
        question = input('input an question:')
        data = handler.classify(question)
        print(data)

